package com.gitee.melin.bee.core.spring.security;

import jakarta.servlet.Filter;
import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.ServletRequest;
import jakarta.servlet.ServletResponse;
import jakarta.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.List;
import java.util.Locale;

/**
 * ref: https://github.com/spring-cloud/spring-cloud-gateway/blob/main/spring-cloud-gateway-server/src/main/java/org/springframework/cloud/gateway/filter/factory/SecureHeadersGatewayFilterFactory.java
 */
public class SecureHeadersFilter implements Filter {

    /**
     * Xss-Protection header name.
     */
    public static final String X_XSS_PROTECTION_HEADER = "X-Xss-Protection";

    /**
     * Strict transport security header name.
     */
    public static final String STRICT_TRANSPORT_SECURITY_HEADER = "Strict-Transport-Security";

    /**
     * Frame options header name.
     */
    public static final String X_FRAME_OPTIONS_HEADER = "X-Frame-Options";

    /**
     * Content-Type Options header name.
     */
    public static final String X_CONTENT_TYPE_OPTIONS_HEADER = "X-Content-Type-Options";

    /**
     * Referrer Policy header name.
     */
    public static final String REFERRER_POLICY_HEADER = "Referrer-Policy";

    /**
     * Content-Security Policy header name.
     */
    public static final String CONTENT_SECURITY_POLICY_HEADER = "Content-Security-Policy";

    /**
     * Download Options header name.
     */
    public static final String X_DOWNLOAD_OPTIONS_HEADER = "X-Download-Options";

    /**
     * Permitted Cross-Domain Policies header name.
     */
    public static final String X_PERMITTED_CROSS_DOMAIN_POLICIES_HEADER = "X-Permitted-Cross-Domain-Policies";

    private SecureHeadersProperties properties;

    public SecureHeadersFilter(SecureHeadersProperties properties) {
        this.properties = properties;
    }

    @Override
    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain chain)
            throws ServletException, IOException {

        chain.doFilter(servletRequest, servletResponse);

        HttpServletResponse response = (HttpServletResponse) servletResponse;
        List<String> disabled = properties.getDisable();

        if (isEnabled(disabled, X_XSS_PROTECTION_HEADER)) {
            response.setHeader(X_XSS_PROTECTION_HEADER, properties.getXssProtectionHeader());
        }

        if (isEnabled(disabled, STRICT_TRANSPORT_SECURITY_HEADER)) {
            response.setHeader(STRICT_TRANSPORT_SECURITY_HEADER, properties.getStrictTransportSecurity());
        }

        if (isEnabled(disabled, X_FRAME_OPTIONS_HEADER)) {
            response.setHeader(X_FRAME_OPTIONS_HEADER, properties.getFrameOptions());
        }

        if (isEnabled(disabled, X_CONTENT_TYPE_OPTIONS_HEADER)) {
            response.setHeader(X_CONTENT_TYPE_OPTIONS_HEADER, properties.getContentTypeOptions());
        }

        if (isEnabled(disabled, REFERRER_POLICY_HEADER)) {
            response.setHeader(REFERRER_POLICY_HEADER, properties.getReferrerPolicy());
        }

        if (isEnabled(disabled, CONTENT_SECURITY_POLICY_HEADER)) {
            response.setHeader(CONTENT_SECURITY_POLICY_HEADER, properties.getContentSecurityPolicy());
        }

        if (isEnabled(disabled, X_DOWNLOAD_OPTIONS_HEADER)) {
            response.setHeader(X_DOWNLOAD_OPTIONS_HEADER, properties.getDownloadOptions());
        }

        if (isEnabled(disabled, X_PERMITTED_CROSS_DOMAIN_POLICIES_HEADER)) {
            response.setHeader(X_PERMITTED_CROSS_DOMAIN_POLICIES_HEADER, properties.getPermittedCrossDomainPolicies());
        }
    }

    private boolean isEnabled(List<String> disabledHeaders, String header) {
        return !disabledHeaders.contains(header.toLowerCase(Locale.ROOT));
    }
}
